// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/common/base/exceptions.h>
#include <pollux/core/expressions.h>
#include <pollux/parse/i_expr.h>
#include <pollux/type/variant.h>

namespace kumo::pollux::core {
    class CallExpr;
    class LambdaExpr;
    class FieldAccessExpr;

    class Expressions {
    public:
        using TypeResolverHook = std::function<TypePtr(
            const std::vector<core::TypedExprPtr> &inputs,
            const std::shared_ptr<const core::CallExpr> &expr,
            bool nullOnFailure)>;

        using FieldAccessHook = std::function<core::TypedExprPtr(
            std::shared_ptr<const FieldAccessExpr> fae,
            std::vector<core::TypedExprPtr> &children)>;

        static TypedExprPtr inferTypes(
            const std::shared_ptr<const IExpr> &expr,
            const TypePtr &input,
            memory::MemoryPool *pool,
            const VectorPtr &complexConstants = nullptr);

        static TypePtr getInputRowType(const TypedExprPtr &expr);

        static void setTypeResolverHook(TypeResolverHook hook) {
            resolverHook_ = hook;
        }

        static TypeResolverHook getResolverHook() {
            return resolverHook_;
        }

        static void setFieldAccessHook(FieldAccessHook hook) {
            fieldAccessHook_ = hook;
        }

        static FieldAccessHook getFieldAccessHook() {
            return fieldAccessHook_;
        }

        static TypedExprPtr inferTypes(
            const std::shared_ptr<const IExpr> &expr,
            const TypePtr &input,
            const std::vector<TypePtr> &lambdaInputTypes,
            memory::MemoryPool *pool,
            const VectorPtr &complexConstants = nullptr);

    private:
        static TypedExprPtr resolveLambdaExpr(
            const std::shared_ptr<const core::LambdaExpr> &lambdaExpr,
            const TypePtr &inputRow,
            const std::vector<TypePtr> &lambdaInputTypes,
            memory::MemoryPool *pool,
            const VectorPtr &complexConstants = nullptr);

        static TypedExprPtr tryResolveCallWithLambdas(
            const std::shared_ptr<const CallExpr> &expr,
            const TypePtr &input,
            memory::MemoryPool *pool,
            const VectorPtr &complexConstants = nullptr);

        static TypeResolverHook resolverHook_;
        static FieldAccessHook fieldAccessHook_;
    };

    class InputExpr : public core::IExpr {
    public:
        InputExpr() {
        }

        std::string toString() const override {
            return "ROW";
        }

        const std::vector<std::shared_ptr<const IExpr> > &getInputs() const override {
            return EMPTY();
        }

        POLLUX_DEFINE_CLASS_NAME(InputExpr)
    };

    class FieldAccessExpr : public core::IExpr {
    public:
        explicit FieldAccessExpr(
            const std::string &name,
            std::optional<std::string> alias,
            std::vector<std::shared_ptr<const IExpr> > &&inputs =
                    std::vector<std::shared_ptr<const IExpr> >{
                        std::make_shared<const InputExpr>()
                    })
            : IExpr{std::move(alias)}, name_{name}, inputs_{std::move(inputs)} {
            KCHECK_EQ(inputs_.size(), 1);
        }

        const std::string &getFieldName() const {
            return name_;
        }

        bool isRootColumn() const {
            if (UNLIKELY(inputs_.empty())) {
                return false;
            }
            return dynamic_cast<const core::InputExpr *>(inputs_.front().get()) !=
                   nullptr;
        }

        std::string toString() const override {
            if (UNLIKELY(inputs_.empty())) {
                return appendAliasIfExists(toStringForRootColumn());
            }
            return appendAliasIfExists(
                isRootColumn() ? toStringForRootColumn() : toStringForMemberAccess());
        }

        const std::vector<std::shared_ptr<const IExpr> > &getInputs() const override {
            return inputs_;
        }

    private:
        std::string toStringForRootColumn() const {
            return "\"" + getEscapedName() + "\"";
        }

        std::string toStringForMemberAccess() const {
            return "dot(" + inputs_.front()->toString() + ",\"" + getEscapedName() +
                   "\")";
        }

        std::string getEscapedName() const {
            return melon::cEscape<std::string>(name_);
        }

        const std::string name_;
        const std::vector<std::shared_ptr<const IExpr> > inputs_;

        POLLUX_DEFINE_CLASS_NAME(FieldAccessExpr)
    };

    class CallExpr : public core::IExpr {
    public:
        CallExpr(
            std::string &&funcName,
            std::vector<std::shared_ptr<const IExpr> > &&inputs,
            std::optional<std::string> alias)
            : core::IExpr{std::move(alias)},
              name_{std::move(funcName)},
              inputs_{std::move(inputs)} {
            POLLUX_CHECK(!name_.empty());
        }

        const std::string &getFunctionName() const {
            return name_;
        }

        std::string toString() const override {
            std::string buf{name_ + "("};
            bool first = true;
            for (auto &f: getInputs()) {
                if (!first) {
                    buf += ",";
                }
                buf += f->toString();
                first = false;
            }
            buf += ")";
            return appendAliasIfExists(buf);
        }

        const std::vector<std::shared_ptr<const IExpr> > &getInputs() const override {
            return inputs_;
        }

    private:
        const std::string name_;
        const std::vector<std::shared_ptr<const IExpr> > inputs_;

        POLLUX_DEFINE_CLASS_NAME(CallExpr)
    };

    class ConstantExpr : public IExpr,
                         public std::enable_shared_from_this<ConstantExpr> {
    public:
        explicit ConstantExpr(
            TypePtr type,
            variant value,
            std::optional<std::string> alias)
            : IExpr{std::move(alias)},
              type_{std::move(type)},
              value_{std::move(value)} {
        }

        std::string toString() const override {
            return appendAliasIfExists(variant{value_}.toJson(type_));
        }

        const variant &value() const {
            return value_;
        }

        const TypePtr &type() const {
            return type_;
        }

        const std::vector<std::shared_ptr<const IExpr> > &getInputs() const override {
            return EMPTY();
        }

        POLLUX_DEFINE_CLASS_NAME(ConstantExpr)

    private:
        const TypePtr type_;
        const variant value_;
    };

    class CastExpr : public IExpr, public std::enable_shared_from_this<CastExpr> {
    private:
        const TypePtr type_;
        const std::shared_ptr<const IExpr> expr_;
        const std::vector<std::shared_ptr<const IExpr> > inputs_;
        bool nullOnFailure_;

    public:
        explicit CastExpr(
            const TypePtr &type,
            const std::shared_ptr<const IExpr> &expr,
            bool nullOnFailure,
            std::optional<std::string> alias)
            : IExpr{std::move(alias)},
              type_(type),
              expr_(expr),
              inputs_({expr}),
              nullOnFailure_(nullOnFailure) {
        }

        std::string toString() const override {
            return appendAliasIfExists(
                "cast(" + expr_->toString() + ", " + type_->toString() + ")");
        }

        const std::vector<std::shared_ptr<const IExpr> > &getInputs() const override {
            return inputs_;
        }

        const TypePtr &type() const {
            return type_;
        }

        const std::shared_ptr<const IExpr> &expr() const {
            return expr_;
        }

        bool nullOnFailure() const {
            return nullOnFailure_;
        }

        POLLUX_DEFINE_CLASS_NAME(CastExpr)
    };

    /// Represents lambda expression as a list of inputs and the body expression.
    /// For example, the expression
    ///     (k, v) -> k + v
    /// is represented using [k, v] as inputNames and k + v as body.
    class LambdaExpr : public IExpr,
                       public std::enable_shared_from_this<LambdaExpr> {
    public:
        LambdaExpr(
            std::vector<std::string> inputNames,
            std::shared_ptr<const IExpr> body)
            : inputNames_{std::move(inputNames)}, body_{{std::move(body)}} {
            POLLUX_CHECK(!inputNames_.empty());
        }

        const std::vector<std::string> &inputNames() const {
            return inputNames_;
        }

        const std::shared_ptr<const IExpr> &body() const {
            return body_[0];
        }

        std::string toString() const override {
            std::ostringstream out;
            if (inputNames_.size() > 1) {
                out << "(";
                for (auto i = 0; i < inputNames_.size(); ++i) {
                    if (i > 0) {
                        out << ", ";
                    }
                    out << inputNames_[i];
                }
                out << ")";
            } else {
                out << inputNames_[0];
            }
            out << " -> " << body_[0]->toString();
            return out.str();
        }

        const std::vector<std::shared_ptr<const IExpr> > &getInputs() const override {
            return body_;
        }

    private:
        std::vector<std::string> inputNames_;
        std::vector<std::shared_ptr<const IExpr> > body_;
    };
} // namespace kumo::pollux::core
